#!/bin/bash
# ==============================================
# 颜色定义和日志函数
# ==============================================

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

# 错误处理函数
handle_error() {
    local message="$1"
    log_error "$message"
    sleep 6
    return 1
}

# 全局配置变量
USE_UV=true  # 默认使用 uv 包管理器

# 获取当前虚拟环境路径
get_current_virtual_env() {
    if [[ -n "$VIRTUAL_ENV" ]]; then
        echo "$VIRTUAL_ENV"
    elif [[ -n "$CONDA_DEFAULT_ENV" ]]; then
        echo "conda:$CONDA_DEFAULT_ENV"
    else
        echo "系统Python环境"
    fi
}

# 包管理器选择函数
choose_package_manager() {
    echo "========================================"
    echo "    包管理器选择"
    echo "========================================"
    echo "请选择要使用的包管理器:"
    echo "1) uv (推荐，更快的Python包管理器)"
    echo "2) pip (传统的Python包管理器)"
    echo "========================================"
    read -p "请选择 (1-2): " choice

    case $choice in
        1)
            USE_UV=true
            log_success "已选择 uv 包管理器"
            ;;
        2)
            USE_UV=false
            log_success "已选择 pip 包管理器"
            ;;
        *)
            log_error "无效选择，默认使用 uv"
            USE_UV=true
            ;;
    esac
}

# ==============================================
# 环境检查和工具函数
# ==============================================

# 检查命令是否存在
check_command() {
    if ! command -v "$1" &> /dev/null; then
        log_error "命令 $1 未找到，请确保已安装"
        return 1
    fi
    return 0
}

# 检查脚本执行方式
check_execution_method() {
    log_info "请使用 source uv-env.sh 来执行此脚本，以确保环境变量生效。"
}

# 检查项目根目录
check_project_root() {
    local readme_found=false
    
    # 遍历当前目录下可能的 README 文件名
    for readme in README README.md ReadMe Readme.md README.MD README.txt; do
        if [[ -f "$readme" ]]; then
            readme_found=true
            break
        fi
    done

    # 若未找到 README，提示用户确认脚本位置
    if [[ "$readme_found" == false ]]; then
        read -p "未在当前目录找到 README 文件，是否已将本脚本放置于项目根目录？(y/n): " confirm
        if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
            handle_error "请将脚本移动到项目根目录后重新运行！"
        fi
    fi
}

# 取消 Conda 环境激活
deactivate_conda() {
    if [[ -n "$CONDA_DEFAULT_ENV" ]]; then
        log_info "检测到当前已激活 Conda 环境: $CONDA_DEFAULT_ENV"
        
        if command -v conda &>/dev/null; then
            log_info "正在取消 Conda 环境激活..."
            conda deactivate
            log_success "Conda 环境已取消激活"
            conda config --set auto_activate_base false
            log_success "已设置 Conda 不自动激活 base 环境"
        else
            log_warn "未找到 conda 命令，但检测到 CONDA_DEFAULT_ENV 环境变量，建议手动取消 Conda 激活后重试"
            read -p "是否继续执行脚本？(y/n): " confirm
            if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
                handle_error "用户取消执行"
            fi
        fi
    fi
}

# ==============================================
# 包管理工具安装函数
# ==============================================

# 安装 pip
install_pip() {
    log_info "正在检测 pip 是否安装..."
    if command -v pip &>/dev/null; then
        log_success "pip 已安装"
        return 0
    fi

    log_warn "未检测到 pip，正在通过 apt 安装..."
    bash <(curl -sSL https://linuxmirrors.cn/main.sh)
    bash <(curl -sSL https://linuxmirrors.cn/docker.sh)
    # 检查是否为 root 用户
    if [[ $EUID -ne 0 ]]; then
        log_info "安装 pip 需要 root 权限，请输入 sudo 密码："
        sudo apt update || handle_error "apt update 失败！"
        sudo apt install -y python3-pip || handle_error "安装 pip 失败！"
    else
        apt update || handle_error "apt update 失败！"
        apt install -y python3-pip || handle_error "安装 pip 失败！"
    fi
    
    log_success "pip 安装成功"
}

# 安装 uv
install_uv() {
    if command -v uv &>/dev/null; then
        log_success "uv 已安装"
        return 0
    fi

    log_info "正在安装 uv..."
    pip install --upgrade pip &> /dev/null || handle_error "pip 升级失败！"
    pip install uv &> /dev/null || handle_error "uv 安装失败！"
    log_success "uv 安装成功"
}

# ==============================================
# Python 版本管理函数
# ==============================================

# 固定 Python 版本
pin_python_version() {
    log_info "正在获取可用的 Python 版本..."
    local available_versions=$(uv python list)

    if [ -z "$available_versions" ]; then
        handle_error "无法获取可用的 Python 版本。"
    fi

    echo -e "ℹ️   所有可使用的 Python 版本列表如下："
    echo "$available_versions"

    echo -e "ℹ️   检测到环境中已安装的 Python 版本有："
    echo -e "     推荐使用以下环境，无需下载"
    echo -e "$(uv python list --only-installed)"

    # 让用户选择版本
    read -p $'ℹ️  请输入您想使用的版本(3.10.12 or 全名)后回车：' selection
    if [ -z "$selection" ]; then
        handle_error "未输入任何版本，操作取消。"
    fi

    uv python pin "$selection" &> /dev/null
    log_success "成功固定 Python 版本为：$selection"
}

# 管理 Python 版本
manage_python_version() {
    if [ -f ".python-version" ]; then
        log_info ".python-version 文件已存在，跳过 Python 版本选择步骤。"
    else
        log_info ".python-version 文件不存在，进入 Python 版本选择步骤。"
        pin_python_version
    fi
}

# ==============================================
# 项目配置函数
# ==============================================

# 配置 pyproject.toml
configure_pyproject() {
    local pyproject_file="pyproject.toml"
    local uv_config_marker="[[tool.uv.index]]"

    if [[ ! -f "$pyproject_file" ]]; then
        log_warn "未找到 $pyproject_file，跳过配置追加步骤"
        uv init || handle_error "uv init 失败！"
        return 0
    fi

    log_info "检测到 $pyproject_file，正在检查 uv 配置..."
    
    if grep -qF "$uv_config_marker" "$pyproject_file"; then
        log_info "$pyproject_file 中已存在 uv 镜像源配置，跳过追加步骤"
        return 0
    fi

    log_info "未检测到 uv 配置，正在追加..."
    
    local uv_config=$(cat << 'EOF'

[[tool.uv.index]]
url = "https://mirrors.cernet.edu.cn/pypi/web/simple"
default = true
marker = "platform_system == 'Linux' and platform_machine=='x86_64'"

[tool.uv.sources]
torch = [
  #{ index = "pytorch-cpu", marker = "platform_system == 'Windows'" },
  { index = "pytorch-cu124", marker = "platform_system == 'Linux' and platform_machine=='x86_64'" },
]
torchvision = [
  #{ index = "pytorch-cpu", marker = "platform_system == 'Windows'" },
  { index = "pytorch-cu124", marker = "platform_system == 'Linux' and platform_machine=='x86_64'" },
]

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[[tool.uv.index]]
name = "pytorch-cu124"
url = "https://download.pytorch.org/whl/cu124"
explicit = true
EOF
    )
    
    echo "$uv_config" >> "$pyproject_file" || handle_error "追加 uv 配置失败！"
    log_success "uv 镜像源配置已成功追加到 $pyproject_file"
}

# 同步项目依赖
sync_project_dependencies() {
    local pyproject_file="pyproject.toml"

    if [[ ! -f "$pyproject_file" ]]; then
        log_warn "未找到 $pyproject_file，跳过 uv sync 和 uv lock"
        return 0
    fi

    log_info "正在运行 uv sync..."
    uv sync || handle_error "uv sync 失败！"
    log_success "uv sync 完成"
    
    log_info "正在运行 uv lock..."
    uv lock || handle_error "uv lock 失败！"
    log_success "uv lock 完成"
}

# ==============================================
# 虚拟环境管理函数
# ==============================================

# 激活虚拟环境
activate_venv() {
    local venv_path=".venv/bin/activate"
    
    if [[ ! -f "$venv_path" ]]; then
        handle_error "未找到虚拟环境激活脚本: $venv_path"
    fi

    log_info "正在激活虚拟环境..."
    source "$venv_path"
    log_success "虚拟环境已激活！当前环境: $(python --version)"
    log_info "提示：脚本执行完毕，虚拟环境仅在当前终端会话有效"
}

# ==============================================
# 依赖包自动安装功能
# ==============================================

# 模块名称映射（特殊模块处理）
declare -A MODULE_MAP=(
    ["cv2"]="opencv-python"
    ["PIL"]="Pillow"
    ["skimage"]="scikit-image"
    ["yaml"]="PyYAML"
    ["MySQLdb"]="mysqlclient"
    ["bs4"]="beautifulsoup4"
    ["distutils"]="setuptools"
)

# 获取实际的pip包名
get_pip_package_name() {
    local module_name="$1"
    
    #此处添加一些错误处理逻辑
    if [[ "$module_name" == "distutils" ]]; then
        log_warn "模块 distutils 通常包含在 setuptools 中，建议安装 setuptools 包,此问题为版本问题"
    fi
    if [[ "$module_name" == "tkinter" ]]; then
        log_warn "模块 tkinter 通常包含在 python-tk 包中，建议安装 python3-tk 包，此问题为系统包问题"
    fi

    if [[ -n "${MODULE_MAP[$module_name]}" ]]; then
        echo "${MODULE_MAP[$module_name]}"
        return 0
    fi
    
    echo "$module_name"
}

# 从Python文件中提取导入的模块
extract_imports_from_file() {
    local file="$1"
    local imports=()
    
    if [[ ! -f "$file" ]]; then
        log_error "文件不存在: $file"
        return 1
    fi
    
    while IFS= read -r line; do
        if [[ "$line" =~ ^[[:space:]]*import[[:space:]]+([a-zA-Z_][a-zA-Z0-9_]*)([[:space:]]|$) ]]; then
            imports+=("${BASH_REMATCH[1]}")
        elif [[ "$line" =~ ^[[:space:]]*from[[:space:]]+([a-zA-Z_][a-zA-Z0-9_]*)[[:space:]]+import ]]; then
            imports+=("${BASH_REMATCH[1]}")
        fi
    done < "$file"
    
    printf '%s\n' "${imports[@]}" | sort -u
}

# 安装Python包
install_python_package() {
    local package="$1"

    log_info "正在安装包: $package"

    if [[ "$USE_UV" == "true" ]]; then
        # 使用 uv 安装
        if uv add "$package"; then
            log_success "成功安装: $package (使用 uv)"
            return 0
        else
            log_error "安装失败: $package (使用 uv)"
            return 1
        fi
    else
        # 使用 pip 安装
        local current_env=$(get_current_virtual_env)
        log_info "当前环境: $current_env"
        log_info "使用 pip 安装包，将安装到当前 Python 环境"

        if pip install "$package"; then
            log_success "成功安装: $package (使用 pip)"
            return 0
        else
            log_error "安装失败: $package (使用 pip)"
            return 1
        fi
    fi
}

# 检查模块是否已安装
is_module_installed() {
    local module="$1"
    python -c "import $module" 2>/dev/null
    return $?
}

# 静态扫描并安装依赖
static_scan_and_install() {
    log_info "开始静态扫描当前目录的Python文件..."
    
    local py_files=()
    while IFS= read -r -d '' file; do
        py_files+=("$file")
    done < <(find . -maxdepth 1 -name "*.py" -type f -print0)
    
    if [[ ${#py_files[@]} -eq 0 ]]; then
        log_warn "当前目录下未找到Python文件"
        return 0
    fi
    
    log_info "找到 ${#py_files[@]} 个Python文件"
    
    local all_imports=()
    for file in "${py_files[@]}"; do
        log_info "扫描文件: $(basename "$file")"
        local file_imports
        file_imports=$(extract_imports_from_file "$file")
        
        while IFS= read -r import; do
            if [[ -n "$import" ]]; then
                all_imports+=("$import")
            fi
        done <<< "$file_imports"
    done
    
    local unique_imports=()
    while IFS= read -r -d '' import; do
        unique_imports+=("$import")
    done < <(printf '%s\n' "${all_imports[@]}" | sort -u | tr '\n' '\0')
    
    if [[ ${#unique_imports[@]} -eq 0 ]]; then
        log_info "未找到需要安装的模块"
        return 0
    fi
    
    log_info "发现 ${#unique_imports[@]} 个唯一模块"
    
    local installed_count=0
    for import in "${unique_imports[@]}"; do
        if python -c "import sys; print('$import' in sys.builtin_module_names)" 2>/dev/null | grep -q "True"; then
            log_info "跳过标准库模块: $import"
            continue
        fi
        
        if ! is_module_installed "$import"; then
            local package_name
            # 获取正确的包名字
            package_name=$(get_pip_package_name "$import")
            
            if install_python_package "$package_name"; then
                ((installed_count++))
            fi
        else
            log_info "模块已安装: $import"
        fi
    done
    
    log_success "静态扫描完成，安装了 $installed_count 个新包"
}

# 执行Python文件并自动安装缺失模块
execute_with_auto_install() {
    local python_file="$1"
    local max_attempts=10
    local attempt=1
    
    if [[ ! -f "$python_file" ]]; then
        log_error "文件不存在: $python_file"
        return 1
    fi
    
    log_info "开始执行: $python_file"
    
    while [[ $attempt -le $max_attempts ]]; do
        log_info "执行尝试: $attempt/$max_attempts"
        
        local output
        output=$(python "$python_file" 2>&1)
        local exit_code=$?
        
        if [[ $exit_code -eq 0 ]]; then
            log_success "Python文件执行成功!"
            echo "$output"
            return 0
        else
            local missing_module=""
            

            # TODO：以下的匹配规则存在问题，现有的包缺失没有匹配上
            # 改进的匹配规则，支持多种错误格式
            missing_module=$(echo "$output" | grep -oE "(ModuleNotFoundError|ImportError): No module named '[^']*'" | cut -d"'" -f2 | cut -d'.' -f1)
            # 如果上面的匹配失败，尝试匹配其他格式
            if [[ -z "$missing_module" ]]; then
                missing_module=$(echo "$output" | grep -o "No module named '[^']*'" | cut -d"'" -f2 | cut -d'.' -f1)
            fi
            if [[ -n "$missing_module" ]]; then
                log_warn "检测到缺失模块: $missing_module"
                
                local package_name
                package_name=$(get_pip_package_name "$missing_module")
                
                if install_python_package "$package_name"; then
                    log_info "等待模块安装完成..."
                    sleep 2
                    ((attempt++))
                    continue
                else
                    log_error "无法安装模块: $missing_module"
                    return 1
                fi
            else
                log_error "执行失败，但不是模块缺失错误:"
                echo "$output"
                return 1
            fi
        fi
    done
    
    log_error "达到最大尝试次数 ($max_attempts)，执行失败"
    return 1
}

# ==============================================
# 用户界面函数
# ==============================================

# 显示模式选择菜单
show_mode_menu() {
    echo "========================================"
    echo "    Python环境自动设置脚本"
    echo "========================================"
    echo "请选择运行模式:"
    echo "1) 环境设置模式 - 设置Python环境和项目依赖"
    echo "2) 静态模式 - 扫描当前目录所有.py文件并安装依赖"
    echo "3) 动态模式 - 执行指定Python文件并自动安装缺失模块"
    echo "4) 退出"
    echo "========================================"
}

# 环境设置模式
environment_setup_mode() {
    log_info "开始环境设置..."

    check_execution_method
    check_project_root
    deactivate_conda
    install_pip

    # 让用户选择包管理器
    choose_package_manager

    if [[ "$USE_UV" == "true" ]]; then
        # uv 模式：安装 uv 并进行相关设置
        install_uv
        manage_python_version
        configure_pyproject
        sync_project_dependencies
        activate_venv

        log_success "环境设置完成！"
        echo -e "\n🎉 脚本执行完成！后续请用 uv run xxx.py 以及安装包 uv add xxx"
    else
        # pip 模式：跳过 uv 相关步骤
        log_info "跳过 uv 相关步骤（Python版本管理、pyproject.toml配置等）"

        # 检查当前环境
        local current_env=$(get_current_virtual_env)
        log_info "当前 Python 环境: $current_env"

        # 询问是否创建虚拟环境
        read -p "是否创建新的虚拟环境？(y/n): " create_venv
        if [[ "$create_venv" =~ ^[Yy]$ ]]; then
            log_info "正在创建虚拟环境..."
            python -m venv .venv || handle_error "创建虚拟环境失败！"
            activate_venv
        else
            log_info "使用现有环境: $current_env"
            log_info "请注意：使用 pip 安装包将安装到当前环境"
        fi

        log_success "环境设置完成！"
        echo -e "\n🎉 脚本执行完成！后续请使用 pip install 安装包"
    fi
}

# 依赖包安装模式选择
dependency_installation_mode() {
    # 显示当前包管理器选择
    if [[ "$USE_UV" == "true" ]]; then
        log_info "当前包管理器: uv"
    else
        log_info "当前包管理器: pip"
    fi

    # 询问是否更改包管理器
    read -p "是否更改包管理器？(y/n): " change_pm
    if [[ "$change_pm" =~ ^[Yy]$ ]]; then
        choose_package_manager
    fi

    while true; do
        echo "========================================"
        echo "    Python环境自动设置脚本"
        echo "========================================"
        echo "当前包管理器: $([[ "$USE_UV" == "true" ]] && echo "uv" || echo "pip")"
        echo "请选择运行模式:"
        echo "1) 静态模式 - 扫描当前目录所有.py文件并安装依赖"
        echo "2) 动态模式 - 执行指定Python文件并自动安装缺失模块"
        echo "3)  退出"
        echo "========================================"
        read -p "请选择模式 (1-3): " choice
        
        case $choice in
            1)
                log_info "选择静态模式（自动扫描py文件进行依赖包安装）"
                static_scan_and_install
                ;;
            2)
                log_info "选择动态模式（执行指定Python文件并自动安装缺失模块）"
                
                # 询问用户要执行的Python文件
                log_info "请输入要执行的Python文件名（在当前目录下）:"
                read -r python_file
                
                # 如果用户没有输入文件名，列出可用的文件
                if [[ -z "$python_file" ]]; then
                    log_info "当前目录下的Python文件:"
                    find . -maxdepth 1 -name "*.py" -type f -exec basename {} \;
                    log_info "请重新输入要执行的Python文件名:"
                    read -r python_file
                fi
                
                if [[ -z "$python_file" ]]; then
                    log_error "未指定Python文件"
                    continue
                fi
                
                # 确保文件存在
                if [[ ! -f "$python_file" ]]; then
                    log_error "文件不存在: $python_file"
                    continue
                fi
                
                # 执行文件并自动安装缺失模块
                execute_with_auto_install "$python_file"
                ;;
            3)
                log_info "退出脚本"
                break
                ;;
            *)
                log_error "无效选择，请输入 1-3"
                ;;
        esac
        
        echo
        read -p "是否继续使用其他模式？(y/n): " continue_choice
        if [[ ! "$continue_choice" =~ ^[Yy]$ ]]; then
            log_info "退出脚本"
            break
        fi
    done
}

# ==============================================
# 主函数
# ==============================================

# 主函数
main() {
    log_info "Python环境自动设置脚本启动"
    local version="1.0.1"
    # 显示欢迎信息
    echo "========================================"
    echo "    Python环境自动设置脚本                 "
    echo "  Author: Aoian       Version: $version"
    echo "请将本脚本放置于项目的readme同级目录运行"
    echo "本脚本采用uv管理python项目，不用 pip 等其他工具"
    echo "========================================"
    echo "请选择运行模式:"
    echo "1) 环境创建（推荐新项目使用）"
    echo "2) 依赖包安装（已有环境，仅管理依赖）"
    echo "3) 退出"
    echo "========================================"
    
    read -p "请选择模式 (1-3): " main_choice
    
    case $main_choice in
        1)
            log_info "环境初始化模式"
            environment_setup_mode
            ;;
        2)
            log_info "依赖包安装模式"
            dependency_installation_mode
            ;;
        3)
            log_info "退出脚本"
            return 0
            ;;
        *)
            log_error "无效选择，请输入 1-3"
            return 1
            ;;
    esac
    
    log_success "脚本执行完成"
}

main